import gradio as gr
import numpy as np
import cv2  # 导入OpenCV库
from pinecone import Pinecone
from collections import Counter

# Initialize Pinecone client
pinecone = Pinecone(api_key="1569e82f-e6cb-45ef-9de2-2cdc5c74809d")
index_name = "mnist-index"

# Connect to the index
index = pinecone.Index(index_name)

def convert_to_grayscale(img_array):
    if img_array.ndim == 3 and img_array.shape[2] == 3:
        grayscale = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
    elif img_array.ndim == 3 and img_array.shape[2] == 4:
        # 只取RGB通道，忽略透明度通道
        img_array = img_array[..., :3]
        grayscale = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
    else:
        grayscale = img_array
    return grayscale

def preprocess_image(img):
    if img.ndim == 3 and img.shape[2] == 4:
        # 只取RGB通道，忽略透明度通道
        img_array = img[..., :3] * (img[..., 3:4] / 255.0)
    else:
        img_array = img

    if img_array.ndim != 2 and img_array.ndim != 3:
        raise ValueError("图像必须是2D（灰度图）或3D（彩色图).")

    img_array = img_array.astype(np.uint8)
    
    img_resized = cv2.resize(img_array, (8, 8), interpolation=cv2.INTER_CUBIC)
    img_array = convert_to_grayscale(img_resized)

    img_array = img_array.reshape(-1) / 255.0
    img_array = img_array.astype(np.float32)

    return img_array

def predict_digit(sketch):
    if sketch.max() == 0:
        return "未检测到输入，请绘制一个数字。"
    
    img_data = sketch
    img = preprocess_image(img_data)
    
    # Debugging: Check if the image vector contains only zeros
    if np.all(img == 0):
        return "图像向量全为零，请检查预处理步骤。"
    
    img = img.reshape(-1) / 255.0  # Ensure the vector is normalized

    # Debugging: Check if the normalized image vector contains only zeros
    if np.all(img == 0):
        return "归一化后的图像向量全为零，请检查预处理步骤。"
    
    # Query Pinecone index for similar vectors
    results = index.query(queries=[img.tolist()], top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    predicted_label = Counter(labels).most_common(1)[0][0]
    
    return str(predicted_label)


# Create Gradio interface
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="在这里绘制你的数字", type="numpy"),
    outputs=gr.Label(label="预测"),
    title="手写数字识别",
    description="在画布上绘制你的数字并获取预测结果。"
)

# Launch Gradio interface
iface.launch()
